Fused ring RS: per-tile bitmask + drain-ring with owner-boundary flus… - #1
Open
OpenDarrenlu wants to merge 1 commit into
Open
Conversation
OpenDarrenlu
force-pushed
the
ltz_gemm_rs
branch
from
July 17, 2026 14:40
85d2c6c to
affe008
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Target: Three performance issues in the cross-rank communication path of the epilogue in the fused BF16 GEMM + push-ring reduce-scatter kernel (
sm100_bf16_gemm_reduce_scatter_ring.cuh):Coarse-grained per-segment flag — Downstream ranks must wait until the upstream rank has stored the entire owner segment before they can begin reducing any tile of that segment. This creates an R-stage body-synchronization ring with zero tile-level overlap.
2-stage
partial_bufpipeline —w3(load) →WG1(add) →WG2(send) share a single 2-stage buffer. This cannot hide cross-rank load/store latency, causing backpressure onWG1→ TMEM fills up → MMA stalls.Per-tile synchronous drain + barrier —
WG2performstma_store_wait<0>()+NamedBarrier::syncon every tile. NVLink round-trip latency sits directly on the pipeline critical path.Measured (4× GB200, tp=4, n=4096 k=2048):
Correctness: Passed for M=256..8192 × 2 iters × 4 ranks (diff ~5e-6, same order of magnitude as baseline).